import sympy as sp
import numpy as np
# pa1,pa2,pa3,p11,p12,p13,p21,p22,p23,p31,p32,p33 = sp.symbols('pa1,pa2,pa3,p11,p12,p13,p21,p22,p23,p31,p32,p33')
# symbol_list = [pa1,pa2,pa3,p11,p12,p13,p21,p22,p23,p31,p32,p33]

# pi = [
#       pa1*p11,
#       pa2*p21+pa2*p22,
#       pa1*p12,
#       pa3*p31+pa3*p33,
#       pa1*p13,
#       pa3*p32,
#       pa2*p23
#       ]

# # 目标函数
# H = 0

# for i in pi:
#     H += -i*sp.log(i,2)

# # 消元
# H = H.subs({pa3:1-pa1-pa2,p13:1-p11-p12,p23:1-p21-p22,p33:1-p31-p32})

# symbol_list.remove(pa3)
# symbol_list.remove(p13)
# symbol_list.remove(p23)
# symbol_list.remove(p33)

# 变量映射转换
# ta1,ta2,t21,t22 = sp.symbols('ta1,ta2,t21,t22')
# H = H.subs({pa1:1/(1+sp.exp(-ta1)),pa2:1/(1+sp.exp(-ta2)),p21:1/(1+sp.exp(-t21)),p22:1/(1+sp.exp(-t22))})
# symbol_list = [ta1,ta2,t21,t22]

# 梯度
def grad(f,symbol_list):
    grad_list = []
    for i in symbol_list:
        grad_list.append(sp.diff(f,i))
    return sp.Matrix(grad_list)

# {变量:数值},绑定
def symbol_value(symbol_list,value_array):
    subs_dict = {}
    for i in range(len(symbol_list)):
        subs_dict[symbol_list[i]] = value_array[i][0]
    return subs_dict

# 梯度上升
def grad_up(goal_func,symbol_list,initial_value_list,step_value=1e-2,max_num=1e8):
    grad_func_mat = grad(goal_func,symbol_list)

    # 起始点
    x = np.array([initial_value_list]).T
    subs_dict = symbol_value(symbol_list,x)
    goal_func_old = -1
    goal_func_new = goal_func.evalf(subs=subs_dict)

    # 变量更新
    count = 0
    while abs(goal_func_new-goal_func_old)>1e-10 and count<max_num:
        goal_func_old = goal_func_new

        x = x + step_value*np.array(grad_func_mat.evalf(subs=subs_dict))
        subs_dict = symbol_value(symbol_list,x)
        goal_func_new = goal_func.evalf(subs=subs_dict)
        count += 1

        # if count >= 5000:
        #     step_value = 0.01
        # if count >= 10000:
        #     step_value = 0.001

        print(x)
        print(goal_func_new)
        print(count)
        print('\n')
    
    if count >= max_num:
        print('超出最大设定迭代次数，收敛失败')
        return -1
    
    return x,goal_func_new

# ans = grad_up(H,symbol_list,[0.1,0.1,0.1,0.1,0.1,0.1,0.1,0.1])
# print(ans)
